package cs276.pe1.spell;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.Map.Entry;
import java.util.logging.Logger;

import cs276.util.Constants;
import cs276.util.Counter;
import cs276.util.IOUtils;
import cs276.util.StringUtils;

/**
 * Basic k-gram spell checker
 * @author alecmgo@gmail.com
 *
 */
public class KGramSpellingCorrector implements SpellingCorrector {
  static final int TOPK_SIZE = 10;
  private static Logger logger = Logger.getLogger("cs276.pe1.spell.KGramSpellingCorrector");
  public Map<String, List<String>> invertedIndex = new HashMap<String, List<String>>();
  public Set<String> wordSet = new HashSet<String>();
  private int kGramSize = 2;

  private String trainingDataFile = Constants.PE1_DIR + "big.txt.gz";

  Counter<String> wordCount = new Counter<String>();
  int numTotalWords = 0;

  /**
   * Initializes spelling corrector by indexing kgrams in words from a file.  
   * Default k-gram size is 2.
   */
  public KGramSpellingCorrector() {
    createIndex(2);
  }

  /**
   * Constructor that accepts k-gram size
   * @param k - size of k-gram
   */
  public KGramSpellingCorrector(int kGramSize) {
    createIndex(kGramSize);
  }

  /**
   * Constructor that accepts custom training data
   * @param trainingDataFile
   */
  public KGramSpellingCorrector(String trainingDataFile) {
    this.trainingDataFile = trainingDataFile;
    createIndex(2);
  }

  /**
   * Constructor for a mock KGramSpellingCorrector (for testing)
   * @param mock
   */
  public KGramSpellingCorrector(boolean mock) {}

  /**
   * Builds a k-gram index
   * @param kGramSize - size of the k-grams
   */
  public void createIndex(int kGramSize) {
    this.kGramSize = kGramSize;
    File path = new File(trainingDataFile);
    for (String line : IOUtils.readLines(IOUtils.openFile(path))) {
      for (String word : StringUtils.tokenize(line)) {
        wordCount.incrementCount(word);
        numTotalWords++;
        if(!wordSet.contains(word)) {
          wordSet.add(word);
          List<String> kgrams = KGramSpellingCorrector.getKGrams(word, kGramSize);
          addToIndex(word, kgrams);
          logger.fine("Adding word: " + word);
        }
      }
    }
  }

  /**
   * Returns a list of possible corrections for a word.  Uses the Jaccard measure.
   */
  public List<String> corrections(String wordToCheck) {
    List<String> kgrams = getKGrams(wordToCheck, kGramSize);
    Counter<String> counter = new Counter<String>();

    //Count how many matches there are
    for(String kgram : kgrams) {
      List<String> wordList = invertedIndex.get(kgram);
      if(wordList != null) {
        for(String word : wordList) {
          counter.incrementCount(word);
        }
      }
    }

    //Apply Jaccard Score by dividing by union
    for(Entry e : counter.entrySet()) {
      String tempWord = (String) e.getKey();
      Double tempCount = (Double) e.getValue();
      double unionSize = wordToCheck.length() + tempWord.length() - tempCount;
      counter.setCount(tempWord, tempCount / unionSize);
    }

    return counter.topK(TOPK_SIZE);
  }

  /**
   * Returns the union size of two words, using letters as the unit
   * @param word1
   * @param word2
   * @return
   */
  public static int getUnionSizeUsingLetters(String word1, String word2) {
    Set<Character> letters = new HashSet<Character>();

    for(int i = 0; i < word1.length(); i++) {
      letters.add(word1.charAt(i));
    }

    for(int i = 0; i < word2.length(); i++) {
      letters.add(word2.charAt(i));
    }

    return letters.size();
  }

  /**
   * Generates a list of k-grams
   * @param word
   * @param k - the size of each k-gram
   * @return
   */
  public static List<String> getKGrams(String word, int k) {
    List<String> kGrams = new ArrayList<String>();

    String start = "$" + word.substring(0, k - 1);
    kGrams.add(start);

    for(int wordIndex = 0; wordIndex < word.length() - k + 1; wordIndex++) {
      kGrams.add(word.substring(wordIndex, wordIndex + k));
    }

    String end = word.substring(word.length() - k + 1) + "$";
    kGrams.add(end);

    return kGrams;
  }

  /**
   * Adds word to inverted k-gram index.
   * @param word
   * @param kgrams - kgrams of the word
   */
  public void addToIndex(String word, List<String> kgrams) {
    for(String kgram : kgrams) {
      List<String> wordList = invertedIndex.get(kgram);
      if(wordList != null) {
        wordList.add(word);
      } else {
        List<String> newWordList = new ArrayList<String>();
        newWordList.add(word);
        invertedIndex.put(kgram, newWordList);
      }
    }
  }

  /**
   * Returns inverted index (for testing)
   * @return
   */
  public Map<String, List<String>> getInvertedIndex() {
    return invertedIndex;
  }
}
